max_iter = 150;
D = 30;

alpha = 1;
U = randn(D,N); 
V = randn(D,M); 
T = randn(D,K);

rmse = zeros(max_iter,1);

for iter=1:max_iter
    %   alpha = sample_alpha(train,U,V,T);
    
    [mu_U, L_U] = sample_user_hyperparams(U);
    
    [mu_V, L_V] = sample_movie_hyperparams(V);
    
    [mu_T, L_T] = sample_time_hyperparams(T);
    
    for ii=1:N
        fil = fU{ii};
        if ~isempty(fil.vals)
            Q = V(:,fil.jj).*T(:,fil.kk);
            QQ = Q*Q';
            QR = Q*fil.vals;
            
            L_i = L_U + alpha*QQ;
            Li_i = inv(L_i);
            mu_i = Li_i*(L_U*mu_U + alpha*QR);
            
            U(:,ii) = mu_i + chol(Li_i)'*randn(D,1);
        end
    end
    
    parfor jj=1:M
        fil = fV{jj};
        if ~isempty(fil.vals)
            P = U(:,fil.ii).*T(:,fil.kk);
            PP = P*P';
            PR = P*fil.vals;
            
            L_j = L_V + alpha*PP;
            Li_j = inv(L_j);
            mu_j = Li_j*(L_V*mu_V + alpha*PR);
            
            V(:,jj) = mu_j + chol(Li_j)'*randn(D,1);
        end
    end
    
    for kk=1:K
        fil = fT{kk};
        if ~isempty(fil.vals)
            X = U(:,fil.ii).*V(:,fil.jj);
            XX = X*X';
            XR = X*fil.vals;
            
            switch kk
                case 1,
                    L_k = 2*L_T + alpha*XX;
                    Li_k = inv(L_k);
                    mu_k = (T(:,2)+mu_T)/2;
                case K,
                    L_k = L_T + alpha*XX;
                    Li_k = inv(L_k);
                    mu_k = Li_k*(L_T*T(:,K-1) + alpha*XR);
                otherwise,
                    L_k = 2*L_T + alpha*XX;
                    Li_k = inv(L_k);
                    mu_k = Li_k*(L_T*(T(:,kk-1)+T(:,kk+1))+ alpha*XR);
            end
            T(:,kk) = mu_k + chol(Li_k)'*randn(D,1);
        end
    end
    
    %if (mod(iter,10)==0) 
    %   z = iter/10;
    fprintf('iter = %d ',iter);
    [out, rmse(iter)] = predict_ratings(train,U,V,T);
    fprintf('rmse = %2.4f\n', rmse(iter))
    %end
    
end